{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# HybridBayesNet"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "license_cell",
      "metadata": {
        "tags": [
          "remove-cell"
        ]
      },
      "source": [
        "GTSAM Copyright 2010-2022, Georgia Tech Research Corporation,\nAtlanta, Georgia 30332-0415\nAll Rights Reserved\n\nAuthors: Frank Dellaert, et al. (see THANKS for the full author list)\n\nSee LICENSE for the license information"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "<a href=\"https://colab.research.google.com/github/borglab/gtsam/blob/develop/gtsam/hybrid/doc/HybridBayesNet.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "tags": [
          "remove-cell"
        ]
      },
      "outputs": [],
      "source": [
        "try:\n",
        "    import google.colab\n",
        "    %pip install --quiet gtsam-develop\n",
        "except ImportError:\n",
        "    pass  # Not running on Colab, do nothing"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "A `HybridBayesNet` represents a directed graphical model (Bayes Net) specifically designed for hybrid systems. It is a collection of `gtsam.HybridConditional` objects, ordered according to an elimination sequence.\n",
        "\n",
        "It extends `gtsam.BayesNet<HybridConditional>` and allows representing the joint probability distribution $P(X, M)$ over continuous variables $X$ and discrete variables $M$ as a product of conditional probabilities:\n",
        "$$\n",
        "P(X, M) = \\prod_i P(\\text{Frontal}_i | \\text{Parents}_i)\n",
        "$$\n",
        "where each conditional $P(\\text{Frontal}_i | \\text{Parents}_i)$ is stored as a `HybridConditional`. This structure allows for representing complex dependencies, such as continuous variables conditioned on discrete modes ($P(X|M)$) alongside purely discrete ($P(M)$) or purely continuous ($P(X)$) relationships.\n",
        "\n",
        "`HybridBayesNet` objects are typically obtained by eliminating a `HybridGaussianFactorGraph`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {},
      "outputs": [],
      "source": [
        "import gtsam\n",
        "import numpy as np\n",
        "import graphviz\n",
        "\n",
        "from gtsam import (\n",
        "    HybridConditional,\n",
        "    GaussianConditional,\n",
        "    DiscreteConditional,\n",
        "    HybridGaussianConditional,\n",
        "    HybridGaussianFactorGraph,\n",
        "    HybridGaussianFactor,\n",
        "    JacobianFactor,\n",
        "    DecisionTreeFactor,\n",
        "    Ordering,\n",
        ")\n",
        "from gtsam.symbol_shorthand import X, D"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Creating a HybridBayesNet\n",
        "\n",
        "While they can be constructed manually by adding `HybridConditional`s, they are more commonly obtained via elimination of a `HybridGaussianFactorGraph`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Manually Constructed HybridBayesNet:\n"
          ]
        },
        {
          "data": {
            "image/svg+xml": [
              "<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"no\"?>\n",
              "<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.1//EN\"\n",
              " \"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\">\n",
              "<!-- Generated by graphviz version 12.0.0 (0)\n",
              " -->\n",
              "<!-- Pages: 1 -->\n",
              "<svg width=\"62pt\" height=\"188pt\"\n",
              " viewBox=\"0.00 0.00 62.00 188.00\" xmlns=\"http://www.w3.org/2000/svg\" xmlns:xlink=\"http://www.w3.org/1999/xlink\">\n",
              "<g id=\"graph0\" class=\"graph\" transform=\"scale(1 1) rotate(0) translate(4 184)\">\n",
              "<polygon fill=\"white\" stroke=\"none\" points=\"-4,4 -4,-184 58,-184 58,4 -4,4\"/>\n",
              "<!-- var7205759403792793600 -->\n",
              "<g id=\"node1\" class=\"node\">\n",
              "<title>var7205759403792793600</title>\n",
              "<ellipse fill=\"none\" stroke=\"black\" cx=\"27\" cy=\"-162\" rx=\"27\" ry=\"18\"/>\n",
              "<text text-anchor=\"middle\" x=\"27\" y=\"-156.95\" font-family=\"Times,serif\" font-size=\"14.00\">d0</text>\n",
              "</g>\n",
              "<!-- var8646911284551352320 -->\n",
              "<g id=\"node2\" class=\"node\">\n",
              "<title>var8646911284551352320</title>\n",
              "<ellipse fill=\"none\" stroke=\"black\" cx=\"27\" cy=\"-90\" rx=\"27\" ry=\"18\"/>\n",
              "<text text-anchor=\"middle\" x=\"27\" y=\"-84.95\" font-family=\"Times,serif\" font-size=\"14.00\">x0</text>\n",
              "</g>\n",
              "<!-- var7205759403792793600&#45;&gt;var8646911284551352320 -->\n",
              "<g id=\"edge2\" class=\"edge\">\n",
              "<title>var7205759403792793600&#45;&gt;var8646911284551352320</title>\n",
              "<path fill=\"none\" stroke=\"black\" d=\"M27,-143.7C27,-136.41 27,-127.73 27,-119.54\"/>\n",
              "<polygon fill=\"black\" stroke=\"black\" points=\"30.5,-119.62 27,-109.62 23.5,-119.62 30.5,-119.62\"/>\n",
              "</g>\n",
              "<!-- var8646911284551352321 -->\n",
              "<g id=\"node3\" class=\"node\">\n",
              "<title>var8646911284551352321</title>\n",
              "<ellipse fill=\"none\" stroke=\"black\" cx=\"27\" cy=\"-18\" rx=\"27\" ry=\"18\"/>\n",
              "<text text-anchor=\"middle\" x=\"27\" y=\"-12.95\" font-family=\"Times,serif\" font-size=\"14.00\">x1</text>\n",
              "</g>\n",
              "<!-- var8646911284551352320&#45;&gt;var8646911284551352321 -->\n",
              "<g id=\"edge1\" class=\"edge\">\n",
              "<title>var8646911284551352320&#45;&gt;var8646911284551352321</title>\n",
              "<path fill=\"none\" stroke=\"black\" d=\"M27,-71.7C27,-64.41 27,-55.73 27,-47.54\"/>\n",
              "<polygon fill=\"black\" stroke=\"black\" points=\"30.5,-47.62 27,-37.62 23.5,-47.62 30.5,-47.62\"/>\n",
              "</g>\n",
              "</g>\n",
              "</svg>\n"
            ],
            "text/plain": [
              "<graphviz.sources.Source at 0x1046dd1c0>"
            ]
          },
          "execution_count": 3,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "# --- Method 1: Manual Construction ---\n",
        "hbn_manual = gtsam.HybridBayesNet()\n",
        "\n",
        "# P(D0)\n",
        "dk0 = (D(0), 2)\n",
        "cond_d0 = DiscreteConditional(dk0, [], \"7/3\") # P(D0=0)=0.7\n",
        "hbn_manual.push_back(HybridConditional(cond_d0))\n",
        "\n",
        "# P(X0 | D0)\n",
        "dk0_parent = (D(0), 2)\n",
        " # Mode 0: P(X0 | D0=0) = N(0, 1)\n",
        "gc0 = GaussianConditional(X(0), np.zeros(1), np.eye(1), gtsam.noiseModel.Unit.Create(1))\n",
        " # Mode 1: P(X0 | D0=1) = N(5, 4)\n",
        "gc1 = GaussianConditional(X(0), np.array([2.5]), np.eye(1)*0.5, gtsam.noiseModel.Isotropic.Sigma(1,2.0))\n",
        "cond_x0_d0 = HybridGaussianConditional(dk0_parent, [gc0, gc1])\n",
        "hbn_manual.push_back(HybridConditional(cond_x0_d0))\n",
        "\n",
        "# P(X1 | X0)\n",
        "cond_x1_x0 = GaussianConditional(X(1), np.array([0.0]), np.eye(1), # d, R=I\n",
        "                             X(0), np.eye(1),                  # Parent X0, S=I\n",
        "                             gtsam.noiseModel.Isotropic.Sigma(1, 1.0)) # N(X1; X0, I)\n",
        "hbn_manual.push_back(HybridConditional(cond_x1_x0))\n",
        "\n",
        "print(\"Manually Constructed HybridBayesNet:\")\n",
        "# hbn_manual.print()\n",
        "graphviz.Source(hbn_manual.dot())"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "id": "f11254b8",
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Original HybridGaussianFactorGraph for Elimination:\n"
          ]
        },
        {
          "data": {
            "image/svg+xml": [
              "<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"no\"?>\n",
              "<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.1//EN\"\n",
              " \"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\">\n",
              "<!-- Generated by graphviz version 12.0.0 (0)\n",
              " -->\n",
              "<!-- Pages: 1 -->\n",
              "<svg width=\"206pt\" height=\"84pt\"\n",
              " viewBox=\"0.00 0.00 206.00 83.60\" xmlns=\"http://www.w3.org/2000/svg\" xmlns:xlink=\"http://www.w3.org/1999/xlink\">\n",
              "<g id=\"graph0\" class=\"graph\" transform=\"scale(1 1) rotate(0) translate(4 79.6)\">\n",
              "<polygon fill=\"white\" stroke=\"none\" points=\"-4,4 -4,-79.6 202,-79.6 202,4 -4,4\"/>\n",
              "<!-- var7205759403792793600 -->\n",
              "<g id=\"node1\" class=\"node\">\n",
              "<title>var7205759403792793600</title>\n",
              "<ellipse fill=\"none\" stroke=\"black\" cx=\"27\" cy=\"-57.6\" rx=\"27\" ry=\"18\"/>\n",
              "<text text-anchor=\"middle\" x=\"27\" y=\"-52.55\" font-family=\"Times,serif\" font-size=\"14.00\">d0</text>\n",
              "</g>\n",
              "<!-- factor0 -->\n",
              "<g id=\"node4\" class=\"node\">\n",
              "<title>factor0</title>\n",
              "<ellipse fill=\"black\" stroke=\"black\" cx=\"27\" cy=\"-1.8\" rx=\"1.8\" ry=\"1.8\"/>\n",
              "</g>\n",
              "<!-- var7205759403792793600&#45;&#45;factor0 -->\n",
              "<g id=\"edge1\" class=\"edge\">\n",
              "<title>var7205759403792793600&#45;&#45;factor0</title>\n",
              "<path fill=\"none\" stroke=\"black\" d=\"M27,-39.28C27,-25.77 27,-8.54 27,-3.96\"/>\n",
              "</g>\n",
              "<!-- factor1 -->\n",
              "<g id=\"node5\" class=\"node\">\n",
              "<title>factor1</title>\n",
              "<ellipse fill=\"black\" stroke=\"black\" cx=\"81\" cy=\"-1.8\" rx=\"1.8\" ry=\"1.8\"/>\n",
              "</g>\n",
              "<!-- var7205759403792793600&#45;&#45;factor1 -->\n",
              "<g id=\"edge3\" class=\"edge\">\n",
              "<title>var7205759403792793600&#45;&#45;factor1</title>\n",
              "<path fill=\"none\" stroke=\"black\" d=\"M41.47,-42.18C55.47,-28.24 75.23,-8.54 79.96,-3.84\"/>\n",
              "</g>\n",
              "<!-- var8646911284551352320 -->\n",
              "<g id=\"node2\" class=\"node\">\n",
              "<title>var8646911284551352320</title>\n",
              "<ellipse fill=\"none\" stroke=\"black\" cx=\"99\" cy=\"-57.6\" rx=\"27\" ry=\"18\"/>\n",
              "<text text-anchor=\"middle\" x=\"99\" y=\"-52.55\" font-family=\"Times,serif\" font-size=\"14.00\">x0</text>\n",
              "</g>\n",
              "<!-- var8646911284551352320&#45;&#45;factor1 -->\n",
              "<g id=\"edge2\" class=\"edge\">\n",
              "<title>var8646911284551352320&#45;&#45;factor1</title>\n",
              "<path fill=\"none\" stroke=\"black\" d=\"M93.3,-39.58C88.77,-26.03 82.94,-8.61 81.39,-3.97\"/>\n",
              "</g>\n",
              "<!-- factor2 -->\n",
              "<g id=\"node6\" class=\"node\">\n",
              "<title>factor2</title>\n",
              "<ellipse fill=\"black\" stroke=\"black\" cx=\"117\" cy=\"-1.8\" rx=\"1.8\" ry=\"1.8\"/>\n",
              "</g>\n",
              "<!-- var8646911284551352320&#45;&#45;factor2 -->\n",
              "<g id=\"edge4\" class=\"edge\">\n",
              "<title>var8646911284551352320&#45;&#45;factor2</title>\n",
              "<path fill=\"none\" stroke=\"black\" d=\"M104.7,-39.58C109.23,-26.03 115.06,-8.61 116.61,-3.97\"/>\n",
              "</g>\n",
              "<!-- var8646911284551352321 -->\n",
              "<g id=\"node3\" class=\"node\">\n",
              "<title>var8646911284551352321</title>\n",
              "<ellipse fill=\"none\" stroke=\"black\" cx=\"171\" cy=\"-57.6\" rx=\"27\" ry=\"18\"/>\n",
              "<text text-anchor=\"middle\" x=\"171\" y=\"-52.55\" font-family=\"Times,serif\" font-size=\"14.00\">x1</text>\n",
              "</g>\n",
              "<!-- var8646911284551352321&#45;&#45;factor2 -->\n",
              "<g id=\"edge5\" class=\"edge\">\n",
              "<title>var8646911284551352321&#45;&#45;factor2</title>\n",
              "<path fill=\"none\" stroke=\"black\" d=\"M156.53,-42.18C142.53,-28.24 122.77,-8.54 118.04,-3.84\"/>\n",
              "</g>\n",
              "</g>\n",
              "</svg>\n"
            ],
            "text/plain": [
              "<graphviz.sources.Source at 0x107142c60>"
            ]
          },
          "execution_count": 4,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "# --- Method 2: From Elimination ---\n",
        "hgfg = HybridGaussianFactorGraph()\n",
        "# P(D0) = 70/30\n",
        "hgfg.push_back(DecisionTreeFactor([dk0], \"0.7 0.3\"))\n",
        "# P(X0|D0) = mixture N(0,1); N(5,4)\n",
        "# Factor version: 0.5*|X0-0|^2/1 + C0 ; 0.5*|X0-5|^2/4 + C1\n",
        "factor_gf0 = JacobianFactor(X(0), np.eye(1), np.zeros(1), gtsam.noiseModel.Isotropic.Sigma(1, 1.0))\n",
        "factor_gf1 = JacobianFactor(X(0), np.eye(1), np.array([5.0]), gtsam.noiseModel.Isotropic.Sigma(1, 2.0))\n",
        "# Store -log(prior) for D0 in the hybrid factor (optional, could keep separate)\n",
        "logP_D0_0 = -np.log(0.7)\n",
        "logP_D0_1 = -np.log(0.3)\n",
        "hgfg.push_back(HybridGaussianFactor(dk0, [(factor_gf0, logP_D0_0), (factor_gf1, logP_D0_1)]))\n",
        "# P(X1|X0) = N(X0, 1)\n",
        "hgfg.push_back(JacobianFactor(X(0), -np.eye(1), X(1), np.eye(1), np.zeros(1), gtsam.noiseModel.Isotropic.Sigma(1, 1.0)))\n",
        "\n",
        "print(\"\\nOriginal HybridGaussianFactorGraph for Elimination:\")\n",
        "# hgfg.print()\n",
        "graphviz.Source(hgfg.dot())"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "id": "07db8de1",
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "HybridBayesNet from Elimination:\n"
          ]
        },
        {
          "data": {
            "image/svg+xml": [
              "<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"no\"?>\n",
              "<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.1//EN\"\n",
              " \"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\">\n",
              "<!-- Generated by graphviz version 12.0.0 (0)\n",
              " -->\n",
              "<!-- Pages: 1 -->\n",
              "<svg width=\"62pt\" height=\"188pt\"\n",
              " viewBox=\"0.00 0.00 62.00 188.00\" xmlns=\"http://www.w3.org/2000/svg\" xmlns:xlink=\"http://www.w3.org/1999/xlink\">\n",
              "<g id=\"graph0\" class=\"graph\" transform=\"scale(1 1) rotate(0) translate(4 184)\">\n",
              "<polygon fill=\"white\" stroke=\"none\" points=\"-4,4 -4,-184 58,-184 58,4 -4,4\"/>\n",
              "<!-- var7205759403792793600 -->\n",
              "<g id=\"node1\" class=\"node\">\n",
              "<title>var7205759403792793600</title>\n",
              "<ellipse fill=\"none\" stroke=\"black\" cx=\"27\" cy=\"-162\" rx=\"27\" ry=\"18\"/>\n",
              "<text text-anchor=\"middle\" x=\"27\" y=\"-156.95\" font-family=\"Times,serif\" font-size=\"14.00\">d0</text>\n",
              "</g>\n",
              "<!-- var8646911284551352320 -->\n",
              "<g id=\"node2\" class=\"node\">\n",
              "<title>var8646911284551352320</title>\n",
              "<ellipse fill=\"none\" stroke=\"black\" cx=\"27\" cy=\"-90\" rx=\"27\" ry=\"18\"/>\n",
              "<text text-anchor=\"middle\" x=\"27\" y=\"-84.95\" font-family=\"Times,serif\" font-size=\"14.00\">x0</text>\n",
              "</g>\n",
              "<!-- var7205759403792793600&#45;&gt;var8646911284551352320 -->\n",
              "<g id=\"edge1\" class=\"edge\">\n",
              "<title>var7205759403792793600&#45;&gt;var8646911284551352320</title>\n",
              "<path fill=\"none\" stroke=\"black\" d=\"M27,-143.7C27,-136.41 27,-127.73 27,-119.54\"/>\n",
              "<polygon fill=\"black\" stroke=\"black\" points=\"30.5,-119.62 27,-109.62 23.5,-119.62 30.5,-119.62\"/>\n",
              "</g>\n",
              "<!-- var8646911284551352321 -->\n",
              "<g id=\"node3\" class=\"node\">\n",
              "<title>var8646911284551352321</title>\n",
              "<ellipse fill=\"none\" stroke=\"black\" cx=\"27\" cy=\"-18\" rx=\"27\" ry=\"18\"/>\n",
              "<text text-anchor=\"middle\" x=\"27\" y=\"-12.95\" font-family=\"Times,serif\" font-size=\"14.00\">x1</text>\n",
              "</g>\n",
              "<!-- var8646911284551352320&#45;&gt;var8646911284551352321 -->\n",
              "<g id=\"edge2\" class=\"edge\">\n",
              "<title>var8646911284551352320&#45;&gt;var8646911284551352321</title>\n",
              "<path fill=\"none\" stroke=\"black\" d=\"M27,-71.7C27,-64.41 27,-55.73 27,-47.54\"/>\n",
              "<polygon fill=\"black\" stroke=\"black\" points=\"30.5,-47.62 27,-37.62 23.5,-47.62 30.5,-47.62\"/>\n",
              "</g>\n",
              "</g>\n",
              "</svg>\n"
            ],
            "text/plain": [
              "<graphviz.sources.Source at 0x10716a120>"
            ]
          },
          "execution_count": 5,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "# Note: Using HybridOrdering(hgfg) is generally recommended: \n",
        "# it returns a Colamd constrained ordering where the discrete keys are\n",
        "# eliminated after the continuous keys.\n",
        "ordering = gtsam.HybridOrdering(hgfg)\n",
        "\n",
        "hbn_elim, _ = hgfg.eliminatePartialSequential(ordering)\n",
        "print(\"\\nHybridBayesNet from Elimination:\")\n",
        "# hbn_elim.print()\n",
        "graphviz.Source(hbn_elim.dot())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Operations on HybridBayesNet\n",
        "\n",
        "`HybridBayesNet` allows evaluating the joint probability, sampling, optimizing (finding the MAP state), and extracting marginal or conditional distributions."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "LogProbability P(X0=0.1, X1=0.2, D0=0): -2.160749387689685\n",
            "Probability P(X0=0.1, X1=0.2, D0=0): 0.11523873018620859\n",
            "\n",
            "Sampled HybridValues:\n",
            "HybridValues: \n",
            "  Continuous: 2 elements\n",
            "  x0: 6.29382\n",
            "  x1: 6.6918\n",
            "  Discrete: (d0, 1)\n",
            "  Nonlinear\n",
            "Values with 0 values:\n",
            "\n",
            "MAP Solution (Optimize):\n",
            "HybridValues: \n",
            "  Continuous: 2 elements\n",
            "  x0: 0\n",
            "  x1: 0\n",
            "  Discrete: (d0, 0)\n",
            "  Nonlinear\n",
            "Values with 0 values:\n",
            "\n",
            "MPE Discrete Assignment:\n",
            "DiscreteValues{7205759403792793600: 0}\n"
          ]
        }
      ],
      "source": [
        "# Use the Bayes Net from elimination for consistency\n",
        "hbn = hbn_elim\n",
        "\n",
        "# --- Evaluation ---\n",
        "values = gtsam.HybridValues()\n",
        "values.insert(D(0), 0)\n",
        "values.insert(X(0), np.array([0.1]))\n",
        "values.insert(X(1), np.array([0.2]))\n",
        "\n",
        "log_prob = hbn.logProbability(values)\n",
        "prob = hbn.evaluate(values) # Same as exp(log_prob)\n",
        "print(f\"\\nLogProbability P(X0=0.1, X1=0.2, D0=0): {log_prob}\")\n",
        "print(f\"Probability P(X0=0.1, X1=0.2, D0=0): {prob}\")\n",
        "\n",
        "# --- Sampling ---\n",
        "full_sample = hbn.sample()\n",
        "print(\"\\nSampled HybridValues:\")\n",
        "full_sample.print()\n",
        "\n",
        "# --- Optimization (Finding MAP state) ---\n",
        "# Computes MPE for discrete, then optimizes continuous given MPE\n",
        "map_solution = hbn.optimize()\n",
        "print(\"\\nMAP Solution (Optimize):\")\n",
        "map_solution.print()\n",
        "\n",
        "# --- MPE (Most Probable Explanation for Discrete Variables) ---\n",
        "mpe_assignment = hbn.mpe()\n",
        "print(\"\\nMPE Discrete Assignment:\")\n",
        "print(mpe_assignment) # Should match discrete part of map_solution"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "id": "d8e3e0ee",
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Optimized Continuous Solution for D0=1:\n",
            "VectorValues: 2 elements\n",
            "  x0: 5\n",
            "  x1: 5\n"
          ]
        }
      ],
      "source": [
        "# --- Optimize Continuous given specific Discrete Assignment ---\n",
        "dv = gtsam.DiscreteValues()\n",
        "dv[D(0)] = 1\n",
        "cont_solution_d0_eq_1 = hbn.optimize(dv)\n",
        "print(\"\\nOptimized Continuous Solution for D0=1:\")\n",
        "cont_solution_d0_eq_1.print()\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "id": "758c1790",
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Discrete Marginal P(M):\n",
            "DiscreteBayesNet\n",
            " \n",
            "size: 1\n",
            "conditional 0:  P( d0 ):\n",
            " f[ (d0,2), ]\n",
            "(d0, 0) | 0.731343   | 0\n",
            "(d0, 1) | 0.268657   | 1\n",
            "number of nnzs: 2\n",
            "\n",
            "\n",
            "Gaussian Conditional P(X | D0=0):\n",
            "\n",
            "size: 2\n",
            "conditional 0:  p(x1 | x0)\n",
            "  R = [ 1 ]\n",
            "  S[x0] = [ -1 ]\n",
            "  d = [ 0 ]\n",
            "  logNormalizationConstant: -0.918939\n",
            "  No noise model\n",
            "conditional 1:  p(x0)\n",
            "  R = [ 1 ]\n",
            "  d = [ 0 ]\n",
            "  mean: 1 elements\n",
            "  x0: 0\n",
            "  logNormalizationConstant: -0.918939\n",
            "  No noise model\n"
          ]
        }
      ],
      "source": [
        "# --- Extract Marginal/Conditional Distributions ---\n",
        "# Get P(M) = P(D0)\n",
        "discrete_marginal_bn = hbn.discreteMarginal()\n",
        "print(\"\\nDiscrete Marginal P(M):\")\n",
        "discrete_marginal_bn.print()\n",
        "\n",
        "# Get P(X | M=m) = P(X0, X1 | D0=0)\n",
        "dv[D(0)] = 0\n",
        "gaussian_conditional_bn = hbn.choose(dv)\n",
        "print(\"\\nGaussian Conditional P(X | D0=0):\")\n",
        "gaussian_conditional_bn.print()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Advanced Operations (`errorTree`, `discretePosterior`, `prune`)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Error Tree (Unnormalized Log Posterior Log P'(M|x)) for x0=0.5, x1=1.0:\n",
            "AlgebraicDecisionTreeKey\n",
            " Choice(d0) \n",
            "AlgebraicDecisionTreeKey\n",
            " 0 Leaf 0.56287232\n",
            "AlgebraicDecisionTreeKey\n",
            " 1 Leaf 4.663718\n",
            "\n",
            "Discrete Posterior Tree P(M|x) for x0=0.5, x1=1.0:\n",
            "AlgebraicDecisionTreeKey\n",
            " Choice(d0) \n",
            "AlgebraicDecisionTreeKey\n",
            " 0 Leaf 0.98371106\n",
            "AlgebraicDecisionTreeKey\n",
            " 1 Leaf 0.016288942\n"
          ]
        }
      ],
      "source": [
        "# --- Error Tree (Log P'(M|x) = log P(x|M) + log P(M)) ---\n",
        "# Evaluate unnormalized log posterior of discrete modes given continuous values\n",
        "cont_values_for_error = gtsam.VectorValues()\n",
        "cont_values_for_error.insert(X(0), np.array([0.5]))\n",
        "cont_values_for_error.insert(X(1), np.array([1.0]))\n",
        "\n",
        "error_tree = hbn.errorTree(cont_values_for_error)\n",
        "print(\"\\nError Tree (Unnormalized Log Posterior Log P'(M|x)) for x0=0.5, x1=1.0:\")\n",
        "error_tree.print()\n",
        "\n",
        "# --- Discrete Posterior P(M|x) ---\n",
        "# Normalized version of exp(-errorTree)\n",
        "posterior_tree = hbn.discretePosterior(cont_values_for_error)\n",
        "print(\"\\nDiscrete Posterior Tree P(M|x) for x0=0.5, x1=1.0:\")\n",
        "posterior_tree.print()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "id": "9bec1c66",
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Pruned HybridBayesNet (max_leaves=1):\n",
            "HybridBayesNet\n",
            " \n",
            "size: 2\n",
            "conditional 0: p(x1 | x0)\n",
            "  R = [ 1 ]\n",
            "  S[x0] = [ -1 ]\n",
            "  d = [ 0 ]\n",
            "  logNormalizationConstant: -0.918939\n",
            "  No noise model\n",
            "conditional 1: p(x0)\n",
            "  R = [ 1 ]\n",
            "  d = [ 0 ]\n",
            "  mean: 1 elements\n",
            "  x0: 0\n",
            "  logNormalizationConstant: -0.918939\n",
            "  No noise model\n"
          ]
        },
        {
          "data": {
            "image/svg+xml": [
              "<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"no\"?>\n",
              "<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.1//EN\"\n",
              " \"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\">\n",
              "<!-- Generated by graphviz version 12.0.0 (0)\n",
              " -->\n",
              "<!-- Pages: 1 -->\n",
              "<svg width=\"62pt\" height=\"116pt\"\n",
              " viewBox=\"0.00 0.00 62.00 116.00\" xmlns=\"http://www.w3.org/2000/svg\" xmlns:xlink=\"http://www.w3.org/1999/xlink\">\n",
              "<g id=\"graph0\" class=\"graph\" transform=\"scale(1 1) rotate(0) translate(4 112)\">\n",
              "<polygon fill=\"white\" stroke=\"none\" points=\"-4,4 -4,-112 58,-112 58,4 -4,4\"/>\n",
              "<!-- var8646911284551352320 -->\n",
              "<g id=\"node1\" class=\"node\">\n",
              "<title>var8646911284551352320</title>\n",
              "<ellipse fill=\"none\" stroke=\"black\" cx=\"27\" cy=\"-90\" rx=\"27\" ry=\"18\"/>\n",
              "<text text-anchor=\"middle\" x=\"27\" y=\"-84.95\" font-family=\"Times,serif\" font-size=\"14.00\">x0</text>\n",
              "</g>\n",
              "<!-- var8646911284551352321 -->\n",
              "<g id=\"node2\" class=\"node\">\n",
              "<title>var8646911284551352321</title>\n",
              "<ellipse fill=\"none\" stroke=\"black\" cx=\"27\" cy=\"-18\" rx=\"27\" ry=\"18\"/>\n",
              "<text text-anchor=\"middle\" x=\"27\" y=\"-12.95\" font-family=\"Times,serif\" font-size=\"14.00\">x1</text>\n",
              "</g>\n",
              "<!-- var8646911284551352320&#45;&gt;var8646911284551352321 -->\n",
              "<g id=\"edge1\" class=\"edge\">\n",
              "<title>var8646911284551352320&#45;&gt;var8646911284551352321</title>\n",
              "<path fill=\"none\" stroke=\"black\" d=\"M27,-71.7C27,-64.41 27,-55.73 27,-47.54\"/>\n",
              "<polygon fill=\"black\" stroke=\"black\" points=\"30.5,-47.62 27,-37.62 23.5,-47.62 30.5,-47.62\"/>\n",
              "</g>\n",
              "</g>\n",
              "</svg>\n"
            ],
            "text/plain": [
              "<graphviz.sources.Source at 0x107123890>"
            ]
          },
          "execution_count": 10,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "# --- Pruning ---\n",
        "# Reduces complexity by removing low-probability discrete branches\n",
        "max_leaves = 1 # Force pruning to the most likely mode\n",
        "pruned_hbn = hbn.prune(max_leaves, marginalThreshold=0.8)\n",
        "\n",
        "print(f\"\\nPruned HybridBayesNet (max_leaves={max_leaves}):\")\n",
        "pruned_hbn.print()\n",
        "# Visualize the pruned Bayes Net\n",
        "graphviz.Source(pruned_hbn.dot())"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "py312",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.12.6"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}
